package com.hb.common.util;

import com.hb.common.exception.DbBackupException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.BufferedReader;
import java.io.File;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.util.Arrays;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;

public class ShellExecutorUtil {

    private static final Logger logger = LoggerFactory.getLogger(ShellExecutorUtil.class);
    private static final ExecutorService communicatorExecutor = Executors.newCachedThreadPool();

    /**
     * 执行shell命令
     *
     * @param command   shell命令
     * @param directory 将执行命令的工作目录。
     */
    public static void execute(String[] command, String directory) throws DbBackupException {
        // 记录异常信息
        AtomicReference<String> failMsg = new AtomicReference<>();
        // 执行命令
        try {
            int exit = ShellExecutorUtil.execute(command, directory, 0L, (msg, process) -> failMsg.set(msg));
            if (exit != 0) {
                throw new DbBackupException(failMsg.get());
            }
        } catch (ShellExecutorUtil.CommandTimeoutException e) {
            throw new DbBackupException(e);
        }
    }

    /**
     * 执行shell命令，返回进程退出值。
     *
     * @param command       shell命令
     * @param directory     将执行命令的工作目录。
     * @param timeout       最大等待时间，时间单位：ms。 null表示等待命令执行完成。
     * @param communicators 通信对象，用于接收来自进程的消息行，以及其他操作。
     * @return 进程退出值
     */
    public static int execute(String[] command, String directory, Long timeout, final Communicator... communicators) throws CommandTimeoutException {
        final ProcessBuilder processBuilder = new ProcessBuilder(command);
        if (directory != null) {
            File workDir = new File(directory);
            if (workDir.exists() && workDir.isDirectory()) {
                processBuilder.directory(workDir);
            }
        }
        processBuilder.redirectErrorStream(true);
        int status = -1;
        try {
            final Process process = processBuilder.start();
            if (communicators != null && communicators.length > 0) {
                communicatorExecutor.submit(() -> {
                    BufferedReader reader = null;
                    try {
                        InputStream inputStream = process.getInputStream();
                        if (inputStream == null) {
                            return;
                        }
                        reader = new BufferedReader(new InputStreamReader(inputStream));
                        String line;
                        while ((line = reader.readLine()) != null) {
                            for (Communicator communicator : communicators) {
                                communicator.onMessage(line, process);
                            }
                        }
                    } catch (Exception e) {
                        logger.warn(e.getMessage(), e);
                    } finally {
                        if (reader != null) {
                            try {
                                reader.close();
                            } catch (Exception e) {
                                logger.warn(e.getMessage(), e);
                            }
                        }
                    }
                });
            }
            if (timeout == null || timeout <= 0) {
                status = process.waitFor();
            } else {
                if (!process.waitFor(timeout, TimeUnit.MILLISECONDS)) {
                    throw new CommandTimeoutException(
                            String.format("执行超时: %s, 命令: %s", timeout, Arrays.toString(command)));
                } else {
                    status = process.exitValue();
                }
            }
            TimeUnit.MILLISECONDS.sleep(100);
        } catch (Exception e) {
            logger.warn(e.getMessage(), e);
            if (e instanceof CommandTimeoutException) {
                throw (CommandTimeoutException) e;
            }
        }
        return status;
    }

    /**
     * 超时异常
     */
    public static class CommandTimeoutException extends Exception {
        public CommandTimeoutException(String message) {
            super(message);
        }
    }

    /**
     * 进程执行期间的 shell 命令通信对象。
     */
    public interface Communicator {
        void onMessage(String message, Process process) throws DbBackupException;
    }
}